﻿using System;
using System.Collections.Generic;
using System.Diagnostics.Contracts;
using System.Reactive;
using System.Threading.Tasks;

namespace LinqToOwin
{
  internal sealed class QueryCap : TaskCompletionSource<Unit>, IObserver<object>
  {
    private readonly object gate = new object();
    private readonly LinqToOwinContext context;
    private readonly List<Task> notifications = new List<Task>();
    private bool handled, completed;
    private int pendingNotificationCount;

    public QueryCap(LinqToOwinContext context)
    {
      Contract.Requires(context != null);

      this.context = context;
    }

    [ContractInvariantMethod]
    [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Performance", "CA1822:MarkMembersAsStatic", Justification = "Required for code contracts.")]
    private void ObjectInvariant()
    {
      Contract.Invariant(gate != null);
      Contract.Invariant(context != null);
      Contract.Invariant(notifications != null);
      Contract.Invariant(pendingNotificationCount >= 0);
    }

    public void OnNext(object value)
    {
      var notified = false;
      try
      {
        lock (gate)
        {
          if (completed)
          {
            return;
          }
          else
          {
            try
            {
              // Finally used for atomic operation in case of thread abort.
            }
            finally
            {
              pendingNotificationCount++;
              notified = true;
            }

            handled = true;
          }
        }

        notifications.Add(Responder.RespondAsync(context, value));
      }
      finally
      {
        if (notified)
        {
          lock (gate)
          {
            pendingNotificationCount--;

            if (completed && pendingNotificationCount == 0)
            {
              EndResponse();
            }
          }
        }
      }
    }

    public void OnError(Exception error)
    {
      lock (gate)
      {
        if (!completed)
        {
          completed = true;

          EndResponse(error);
        }
      }
    }

    public void OnCompleted()
    {
      lock (gate)
      {
        if (!completed)
        {
          completed = true;

          if (pendingNotificationCount == 0)
          {
            EndResponse();
          }
        }
      }
    }

    private void EndResponse(Exception error = null)
    {
      if (context.SetHandledUponNotification)
      {
        if (handled)
        {
          context.Handled = true;
        }
      }

      System.Threading.Tasks.Task.WhenAll(notifications).ContinueWith(
        task =>
        {
          if (error != null)
          {
            if (task.IsFaulted)
            {
              TrySetException(new AggregateException(error, task.Exception).Flatten());
            }
            else
            {
              TrySetException(error);
            }
          }
          else if (task.IsFaulted)
          {
            TrySetException(task.Exception);
          }
          else
          {
            TrySetResult(Unit.Default);
          }
        });
    }
  }
}